#!/bin/bash

# Linear Layouts CUDA Implementation - Comprehensive Benchmark Runner
# This script runs all benchmarks and generates performance reports

set -e  # Exit on any error

# Colors for output formatting
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
BUILD_DIR="$PROJECT_DIR/build"
RESULTS_DIR="$PROJECT_DIR/benchmark_results"
TIMESTAMP=$(date "+%Y%m%d_%H%M%S")

# Default values
RUN_TESTS=false
RUN_BENCHMARKS=true
GENERATE_REPORT=true
VERBOSE=false
PROFILE=false

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --help|-h)
            echo "Linear Layouts Benchmark Runner"
            echo "================================"
            echo ""
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --tests, -t           Run tests before benchmarks"
            echo "  --no-benchmarks      Skip benchmark execution"
            echo "  --no-report          Skip report generation"
            echo "  --verbose, -v        Enable verbose output"
            echo "  --profile, -p        Enable profiling with nsys"
            echo "  --help, -h           Show this help message"
            echo ""
            echo "Examples:"
            echo "  $0                   Run all benchmarks (default)"
            echo "  $0 --tests           Run tests first, then benchmarks"
            echo "  $0 --profile         Run with profiling enabled"
            echo "  $0 --verbose         Run with detailed output"
            exit 0
            ;;
        --tests|-t)
            RUN_TESTS=true
            shift
            ;;
        --no-benchmarks)
            RUN_BENCHMARKS=false
            shift
            ;;
        --no-report)
            GENERATE_REPORT=false
            shift
            ;;
        --verbose|-v)
            VERBOSE=true
            shift
            ;;
        --profile|-p)
            PROFILE=true
            shift
            ;;
        *)
            echo "Unknown option: $1"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check prerequisites
check_prerequisites() {
    log_info "Checking prerequisites..."
    
    # Check if we're in the right directory
    if [[ ! -f "$PROJECT_DIR/Makefile" ]]; then
        log_error "Makefile not found. Please run from the project root directory."
        exit 1
    fi
    
    # Check CUDA installation
    if ! command -v nvcc &> /dev/null; then
        log_error "nvcc not found. Please install CUDA toolkit."
        exit 1
    fi
    
    # Check for compatible GPU
    if command -v nvidia-smi &> /dev/null; then
        if ! nvidia-smi --query-gpu=compute_cap --format=csv,noheader,nounits | awk '$1 >= 8.0' | grep -q .; then
            log_warning "No SM_80+ GPU found. Benchmarks may not run optimally."
        fi
    else
        log_warning "nvidia-smi not found. Cannot verify GPU compatibility."
    fi
    
    # Check for profiling tools if requested
    if [[ "$PROFILE" == "true" ]]; then
        if ! command -v nsys &> /dev/null; then
            log_warning "nsys not found. Profiling will be skipped."
            PROFILE=false
        fi
    fi
    
    log_success "Prerequisites check completed"
}

# Create results directory
setup_results_dir() {
    log_info "Setting up results directory..."
    mkdir -p "$RESULTS_DIR"
    
    # Create timestamped subdirectory for this run
    RUN_DIR="$RESULTS_DIR/run_$TIMESTAMP"
    mkdir -p "$RUN_DIR"
    
    # Create subdirectories for different types of results
    mkdir -p "$RUN_DIR/logs"
    mkdir -p "$RUN_DIR/profiles" 
    mkdir -p "$RUN_DIR/reports"
    
    log_success "Results directory created: $RUN_DIR"
}

# Build all targets
build_targets() {
    log_info "Building all targets..."
    
    cd "$PROJECT_DIR"
    
    # Build tests if requested
    if [[ "$RUN_TESTS" == "true" ]]; then
        log_info "Building tests..."
        if [[ "$VERBOSE" == "true" ]]; then
            make tests
        else
            make tests > "$RUN_DIR/logs/build_tests.log" 2>&1
        fi
        log_success "Tests built successfully"
    fi
    
    # Build benchmarks if requested
    if [[ "$RUN_BENCHMARKS" == "true" ]]; then
        log_info "Building benchmarks..."
        if [[ "$VERBOSE" == "true" ]]; then
            make benchmarks
        else
            make benchmarks > "$RUN_DIR/logs/build_benchmarks.log" 2>&1
        fi
        log_success "Benchmarks built successfully"
    fi
}

# Run tests
run_tests() {
    if [[ "$RUN_TESTS" != "true" ]]; then
        return 0
    fi
    
    log_info "Running tests..."
    
    cd "$PROJECT_DIR"
    
    local test_log="$RUN_DIR/logs/tests.log"
    
    # Run each test individually to capture results
    local tests=(
        "test_binary_matrix"
        "test_layouts" 
        "test_operations"
    )
    
    local failed_tests=()
    
    for test in "${tests[@]}"; do
        log_info "Running $test..."
        
        if [[ "$VERBOSE" == "true" ]]; then
            if "./build/$test"; then
                log_success "$test passed"
            else
                log_error "$test failed"
                failed_tests+=("$test")
            fi
        else
            if "./build/$test" >> "$test_log" 2>&1; then
                log_success "$test passed"
            else
                log_error "$test failed (see $test_log for details)"
                failed_tests+=("$test")
            fi
        fi
    done
    
    if [[ ${#failed_tests[@]} -gt 0 ]]; then
        log_error "The following tests failed: ${failed_tests[*]}"
        log_warning "Continuing with benchmarks, but results may be unreliable"
    else
        log_success "All tests passed"
    fi
}

# Run benchmarks
run_benchmarks() {
    if [[ "$RUN_BENCHMARKS" != "true" ]]; then
        return 0
    fi
    
    log_info "Running benchmarks..."
    
    cd "$PROJECT_DIR"
    
    # Get system information
    local system_info="$RUN_DIR/reports/system_info.txt"
    {
        echo "=== System Information ==="
        echo "Date: $(date)"
        echo "Hostname: $(hostname)"
        echo "OS: $(uname -a)"
        echo ""
        echo "=== CUDA Information ==="
        nvcc --version
        echo ""
        echo "=== GPU Information ==="
        nvidia-smi
        echo ""
        echo "=== Memory Information ==="
        free -h
        echo ""
        echo "=== CPU Information ==="
        lscpu
        echo ""
    } > "$system_info"
    
    # Run each benchmark
    local benchmarks=(
        "bench_conversion"
        "bench_vectorization"
        "bench_swizzling"
    )
    
    local failed_benchmarks=()
    
    for benchmark in "${benchmarks[@]}"; do
        log_info "Running $benchmark..."
        
        local bench_log="$RUN_DIR/logs/${benchmark}.log"
        local bench_results="$RUN_DIR/reports/${benchmark}_results.txt"
        
        if [[ "$PROFILE" == "true" ]]; then
            log_info "Profiling $benchmark with nsys..."
            local profile_output="$RUN_DIR/profiles/${benchmark}.nsys-rep"
            
            if nsys profile --output="$profile_output" "./build/$benchmark" > "$bench_results" 2> "$bench_log"; then
                log_success "$benchmark completed (with profiling)"
            else
                log_error "$benchmark failed"
                failed_benchmarks+=("$benchmark")
            fi
        else
            if [[ "$VERBOSE" == "true" ]]; then
                if "./build/$benchmark" | tee "$bench_results"; then
                    log_success "$benchmark completed"
                else
                    log_error "$benchmark failed"
                    failed_benchmarks+=("$benchmark")
                fi
            else
                if "./build/$benchmark" > "$bench_results" 2> "$bench_log"; then
                    log_success "$benchmark completed"
                else
                    log_error "$benchmark failed (see $bench_log for details)"
                    failed_benchmarks+=("$benchmark")
                fi
            fi
        fi
    done
    
    if [[ ${#failed_benchmarks[@]} -gt 0 ]]; then
        log_error "The following benchmarks failed: ${failed_benchmarks[*]}"
        return 1
    else
        log_success "All benchmarks completed successfully"
        return 0
    fi
}

# Generate comprehensive report
generate_report() {
    if [[ "$GENERATE_REPORT" != "true" ]]; then
        return 0
    fi
    
    log_info "Generating comprehensive performance report..."
    
    local report_file="$RUN_DIR/reports/comprehensive_report.md"
    
    {
        echo "# Linear Layouts CUDA Implementation - Benchmark Report"
        echo "Generated on: $(date)"
        echo ""
        
        # System information summary
        echo "## System Configuration"
        echo ""
        echo "- **Date**: $(date)"
        echo "- **Hostname**: $(hostname)"
        echo "- **OS**: $(uname -s -r)"
        echo ""
        
        # GPU information
        if command -v nvidia-smi &> /dev/null; then
            echo "### GPU Configuration"
            echo '```'
            nvidia-smi --query-gpu=name,compute_cap,memory.total --format=csv,noheader
            echo '```'
            echo ""
        fi
        
        # CUDA version
        echo "### CUDA Version"
        echo '```'
        nvcc --version | grep "release"
        echo '```'
        echo ""
        
        # Test results
        if [[ "$RUN_TESTS" == "true" ]]; then
            echo "## Test Results"
            echo ""
            if [[ -f "$RUN_DIR/logs/tests.log" ]]; then
                echo "### Test Summary"
                if grep -q "FAILED" "$RUN_DIR/logs/tests.log"; then
                    echo "❌ Some tests failed. See detailed logs for more information."
                else
                    echo "✅ All tests passed successfully."
                fi
                echo ""
            fi
        fi
        
        # Benchmark results
        echo "## Benchmark Results"
        echo ""
        
        for benchmark in bench_conversion bench_vectorization bench_swizzling; do
            local results_file="$RUN_DIR/reports/${benchmark}_results.txt"
            if [[ -f "$results_file" ]]; then
                echo "### $(echo "$benchmark" | sed 's/_/ /g' | sed 's/\b\w/\U&/g')"
                echo ""
                echo '```'
                cat "$results_file"
                echo '```'
                echo ""
            fi
        done
        
        # Performance analysis
        echo "## Performance Analysis"
        echo ""
        echo "### Key Metrics"
        echo ""
        
        # Extract key metrics from benchmark results
        if [[ -f "$RUN_DIR/reports/bench_conversion_results.txt" ]]; then
            local max_throughput=$(grep -o "[0-9]\+\.[0-9]\+ M ops/s" "$RUN_DIR/reports/bench_conversion_results.txt" | sort -n | tail -1 || echo "N/A")
            echo "- **Maximum Conversion Throughput**: $max_throughput"
        fi
        
        if [[ -f "$RUN_DIR/reports/bench_vectorization_results.txt" ]]; then
            local max_bandwidth=$(grep -o "[0-9]\+\.[0-9]\+ GB/s" "$RUN_DIR/reports/bench_vectorization_results.txt" | sort -n | tail -1 || echo "N/A")
            echo "- **Maximum Memory Bandwidth**: $max_bandwidth"
        fi
        
        if [[ -f "$RUN_DIR/reports/bench_swizzling_results.txt" ]]; then
            local max_speedup=$(grep -o "Speedup: [0-9]\+\.[0-9]\+x" "$RUN_DIR/reports/bench_swizzling_results.txt" | sed 's/Speedup: //' | sed 's/x//' | sort -n | tail -1 || echo "N/A")
            echo "- **Maximum Swizzling Speedup**: ${max_speedup}x"
        fi
        
        echo ""
        echo "### Comparison with Paper Claims"
        echo ""
        echo "The Linear Layouts paper claims the following improvements:"
        echo "- 1.5-14× speedups in micro-benchmarks"
        echo "- 1.1-1.3× improvements in real kernels"
        echo "- Significant reduction in bank conflicts"
        echo ""
        echo "Compare these claims with the actual benchmark results above."
        echo ""
        
        # Profiling information
        if [[ "$PROFILE" == "true" ]]; then
            echo "## Profiling Results"
            echo ""
            echo "Profiling data has been collected using NVIDIA Nsight Systems."
            echo "Profile files are located in the \`profiles/\` directory:"
            echo ""
            for profile in "$RUN_DIR/profiles"/*.nsys-rep; do
                if [[ -f "$profile" ]]; then
                    echo "- $(basename "$profile")"
                fi
            done
            echo ""
            echo "To view the profiles, use: \`nsys-ui <profile_file>\`"
            echo ""
        fi
        
        # Recommendations
        echo "## Recommendations"
        echo ""
        echo "1. **Performance Validation**: Compare benchmark results with paper claims"
        echo "2. **Hardware Utilization**: Check if GPU memory bandwidth is being saturated"
        echo "3. **Layout Selection**: Use appropriate layout types for different workloads"
        echo "4. **Further Optimization**: Consider combining multiple layout optimizations"
        echo ""
        
        # Files generated
        echo "## Generated Files"
        echo ""
        echo "- **System Info**: \`system_info.txt\`"
        echo "- **Build Logs**: \`logs/build_*.log\`"
        echo "- **Test Logs**: \`logs/tests.log\`"
        echo "- **Benchmark Logs**: \`logs/bench_*.log\`"
        echo "- **Benchmark Results**: \`reports/bench_*_results.txt\`"
        if [[ "$PROFILE" == "true" ]]; then
            echo "- **Profile Data**: \`profiles/*.nsys-rep\`"
        fi
        echo ""
        
    } > "$report_file"
    
    log_success "Comprehensive report generated: $report_file"
    
    # Display summary
    echo ""
    echo "=== BENCHMARK SUMMARY ==="
    if [[ -f "$report_file" ]]; then
        echo "Full report: $report_file"
        echo ""
        echo "Key Results:"
        grep -E "Maximum|Speedup" "$report_file" | sed 's/^/  /'
    fi
}

# Main execution
main() {
    echo "=== Linear Layouts CUDA Implementation - Benchmark Runner ==="
    echo ""
    
    # Setup
    check_prerequisites
    setup_results_dir
    
    # Build
    build_targets
    
    # Execute
    run_tests
    local benchmark_success=true
    if ! run_benchmarks; then
        benchmark_success=false
    fi
    
    # Report
    generate_report
    
    # Final summary
    echo ""
    echo "=== EXECUTION SUMMARY ==="
    if [[ "$benchmark_success" == "true" ]]; then
        log_success "All benchmarks completed successfully"
    else
        log_warning "Some benchmarks failed - check logs for details"
    fi
    
    log_info "Results saved to: $RUN_DIR"
    
    if [[ "$GENERATE_REPORT" == "true" ]]; then
        log_info "Comprehensive report: $RUN_DIR/reports/comprehensive_report.md"
    fi
    
    echo ""
    echo "To view results:"
    echo "  cat $RUN_DIR/reports/comprehensive_report.md"
    echo ""
    
    if [[ "$benchmark_success" == "true" ]]; then
        exit 0
    else
        exit 1
    fi
}

# Execute main function
main "$@"